"""
MongoDB索引初始化
为热点数据创建必要的索引以优化查询性能
"""
import asyncio
import os
from motor.motor_asyncio import AsyncIOMotorClient
import logging

logger = logging.getLogger(__name__)

async def create_indexes():
    """创建MongoDB索引"""
    # 连接MongoDB
    client = AsyncIOMotorClient(
        os.getenv('MONGODB_URI', 'mongodb://localhost:27017')
    )
    db = client[os.getenv('MONGODB_DB', 'crawler_db')]
    
    # 支持的平台列表
    platforms = [
        'weibo', 'zhihu', 'toutiao', 'baidu',
        'xiaohongshu', 'douyin', 'twitter', 'reddit', 'youtube'
    ]
    
    created_indexes = []
    
    for platform in platforms:
        collection = db[f'{platform}_hot_items']
        
        try:
            # 创建文本索引用于搜索
            await collection.create_index([
                ('title', 'text'),
                ('excerpt', 'text'),
                ('content', 'text')
            ], name=f'{platform}_text_index')
            
            # 创建复合索引用于查询优化
            await collection.create_index([
                ('crawled_at', -1),
                ('rank', 1)
            ], name=f'{platform}_time_rank_index')
            
            # 创建平台和时间的复合索引
            await collection.create_index([
                ('platform', 1),
                ('crawled_at', -1)
            ], name=f'{platform}_platform_time_index')
            
            # 创建热度值索引
            await collection.create_index([
                ('heat_value', -1)
            ], name=f'{platform}_heat_index')
            
            # 创建唯一索引防止重复
            await collection.create_index([
                ('platform', 1),
                ('item_id', 1),
                ('crawled_at', 1)
            ], unique=True, name=f'{platform}_unique_index')
            
            created_indexes.append(platform)
            logger.info(f"Created indexes for {platform}_hot_items")
            
        except Exception as e:
            logger.error(f"Failed to create indexes for {platform}: {e}")
    
    logger.info(f"Index creation completed. Created indexes for: {created_indexes}")
    return created_indexes

async def verify_indexes():
    """验证索引是否存在"""
    client = AsyncIOMotorClient(
        os.getenv('MONGODB_URI', 'mongodb://localhost:27017')
    )
    db = client[os.getenv('MONGODB_DB', 'crawler_db')]
    
    platforms = [
        'weibo', 'zhihu', 'toutiao', 'baidu',
        'xiaohongshu', 'douyin', 'twitter', 'reddit', 'youtube'
    ]
    
    index_status = {}
    
    for platform in platforms:
        collection = db[f'{platform}_hot_items']
        indexes = await collection.list_indexes().to_list(None)
        index_names = [idx['name'] for idx in indexes]
        
        index_status[platform] = {
            'text_index': f'{platform}_text_index' in index_names,
            'time_rank_index': f'{platform}_time_rank_index' in index_names,
            'platform_time_index': f'{platform}_platform_time_index' in index_names,
            'heat_index': f'{platform}_heat_index' in index_names,
            'unique_index': f'{platform}_unique_index' in index_names
        }
    
    return index_status

# 应用启动时自动创建索引
async def init_indexes_on_startup():
    """在应用启动时初始化索引"""
    try:
        await create_indexes()
        logger.info("MongoDB indexes initialized successfully")
    except Exception as e:
        logger.error(f"Failed to initialize MongoDB indexes: {e}")

if __name__ == "__main__":
    # 命令行执行索引创建
    asyncio.run(create_indexes())